import rclpy
from rclpy.node import Node
import yaml
import math
from geometry_msgs.msg import PointStamped, QuaternionStamped, Quaternion, Point
from nav_msgs.msg import Odometry
from std_msgs.msg import Header
from tf_transformations import quaternion_from_euler

class WaypointController(Node):
    def __init__(self):
        super().__init__('waypoint_controller')
        
        # 读取配置文件
        self.load_config()
        
        # 发布器
        self.attitude_pub_ = self.create_publisher(QuaternionStamped, '/bluerov00/attitude_target_desired', 10)
        self.position_pub_ = self.create_publisher(PointStamped, '/bluerov00/position_target', 10)
        
        # 订阅器
        self.odometry_sub_ = self.create_subscription(
            Odometry, 
            '/bluerov00/odometry', 
            self.odometry_callback, 
            10
        )
        
        # 状态变量
        self.current_waypoint_index = 0
        self.current_position = None
        self.waypoint_reached = False
        self.position_tolerance = 0.1  # 位置容差，单位：米
        
        # 发布第一个目标点
        self.publish_current_waypoint()
        
        self.get_logger().info('Waypoint controller initialized with {} waypoints'.format(len(self.waypoints)))
    
    def load_config(self):
        """加载配置文件"""
        # config = {
        #     'loop': False,
        #     'waypoints': [
        #         {'x': 5.0, 'y': 0.0, 'z': -0.7, 'yaw': 30, 'pitch': 0, 'roll': 0},
        #         {'x': 3.5355339059327378, 'y': 3.5355339059327373, 'z': -0.7, 'yaw': 30, 'pitch': 0, 'roll': 0},
        #         {'x': 0.0, 'y': 5.0, 'z': -0.7, 'yaw': 30, 'pitch': 0, 'roll': 0},
        #         {'x': -3.5355339059327373, 'y': 3.5355339059327378, 'z': -0.7, 'yaw': 30, 'pitch': 0, 'roll': 0}
        #     ]
        # }
        
        config_path = '/home/siyuan/ros2/src/hippo_control/scripts/bernoulli_default.yaml'
        with open(config_path, 'r') as file:
            config = yaml.safe_load(file)
        
        self.loop = config['loop']
        self.waypoints = config['waypoints']
    
    def euler_to_quaternion(self, roll, pitch, yaw):
        """将欧拉角转换为四元数"""
        # 将角度转换为弧度
        roll_rad = math.radians(roll)
        pitch_rad = math.radians(pitch)
        yaw_rad = math.radians(yaw)
        
        # 计算四元数
        q = quaternion_from_euler(roll_rad, pitch_rad, yaw_rad)
        return Quaternion(x=q[0], y=q[1], z=q[2], w=q[3])
    
    def publish_current_waypoint(self):
        """发布当前目标点的位置和姿态"""
        if self.current_waypoint_index < len(self.waypoints):
            waypoint = self.waypoints[self.current_waypoint_index]
            
            # 发布位置
            position_msg = PointStamped()
            position_msg.header = Header()
            position_msg.header.stamp = self.get_clock().now().to_msg()
            position_msg.header.frame_id = 'world'
            position_msg.point = Point(x=waypoint['x'], y=waypoint['y'], z=waypoint['z'])
            self.position_pub_.publish(position_msg)
            
            # 发布姿态
            attitude_msg = QuaternionStamped()
            attitude_msg.header = Header()
            attitude_msg.header.stamp = self.get_clock().now().to_msg()
            attitude_msg.header.frame_id = 'world'
            attitude_msg.quaternion = self.euler_to_quaternion(
                waypoint['roll'], waypoint['pitch'], waypoint['yaw']
            )
            self.attitude_pub_.publish(attitude_msg)
            
            self.get_logger().info('Published waypoint {}: position ({}, {}, {}), attitude (yaw: {}°)'.format(
                self.current_waypoint_index + 1,
                waypoint['x'], waypoint['y'], waypoint['z'],
                waypoint['yaw']
            ))
    
    def is_position_reached(self, current_pos, target_pos):
        """检查是否到达目标位置"""
        if current_pos is None:
            return False
            
        distance = math.sqrt(
            (current_pos.x - target_pos.x) ** 2 +
            (current_pos.y - target_pos.y) ** 2 +
            (current_pos.z - target_pos.z) ** 2
        )
        
        return distance < self.position_tolerance
    
    def odometry_callback(self, msg):
        """里程计回调函数"""
        # 更新当前位置
        self.current_position = msg.pose.pose.position
        
        # 检查是否有目标点需要处理
        if self.current_waypoint_index < len(self.waypoints):
            current_waypoint = self.waypoints[self.current_waypoint_index]
            target_position = Point(
                x=current_waypoint['x'],
                y=current_waypoint['y'],
                z=current_waypoint['z']
            )
            
            # 检查是否到达当前目标点
            if self.is_position_reached(self.current_position, target_position):
                if not self.waypoint_reached:
                    self.get_logger().info('Waypoint {} reached!'.format(self.current_waypoint_index + 1))
                    self.waypoint_reached = True
                    
                    # 移动到下一个目标点
                    self.current_waypoint_index += 1
                    
                    if self.current_waypoint_index < len(self.waypoints):
                        # 发布下一个目标点
                        self.publish_current_waypoint()
                        self.waypoint_reached = False
                    else:
                        self.get_logger().info('All waypoints completed!')
                        
                        # 如果启用循环，重新开始
                        if self.loop:
                            self.current_waypoint_index = 0
                            self.publish_current_waypoint()
                            self.waypoint_reached = False
                            self.get_logger().info('Restarting waypoint sequence...')
            else:
                self.waypoint_reached = False

def main(args=None):
    rclpy.init(args=args)
    waypoint_controller = WaypointController()
    
    try:
        rclpy.spin(waypoint_controller)
    except KeyboardInterrupt:
        pass
    finally:
        waypoint_controller.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()